import queue  
import msgpack  
import time  
from websocket import create_connection, WebSocketConnectionClosedException 
import threading  
from cryptography.fernet import Fernet

# 生成密钥
# 生成密钥并打印
# key = Fernet.generate_key()
# print(f"Generated key: {key.decode()}")  # 将密钥打印出来以便使用
key = b"HwW4y2ZBiCuAWwpdn0arafUIQ4tHlvQql7QvBMNms-w="
cipher = Fernet(key)
def encrypt_data(data):
    """加密数据"""
    packed_data = msgpack.packb(data)
    return cipher.encrypt(packed_data)

def decrypt_data(encrypted_data):
    """解密数据"""
    decrypted_message = cipher.decrypt(encrypted_data)
    return msgpack.unpackb(decrypted_message)


class MapClient:  
    def __init__(self, key):  
        self.url = f"ws://202.189.9.39:8989" #/?key={key}
        self.ws = None  
        self.pending_requests = {}
        self.response_data = {}  # 存储响应数据  
        self._lock = threading.Lock()
        self.request_id = 1
        self.reconnect_interval = 5  # 重连间隔时间，单位秒
        self.heartbeat_interval = 15  # 心跳消息间隔时间，单位秒
        self.heartbeat_fail_threshold = 3  # 心跳发送失败阈值
        self.connect()
        self.start_listening()
        self.start_heartbeat()
  
    def connect(self):  
        while True:  
            try:  
                self.ws = create_connection(self.url)  
                print(f"WebSocket connected to: {self.url}")  
                break  
            except Exception as e:  
                print(f"Failed to connect to WebSocket: {e}")  
                time.sleep(self.reconnect_interval)  

  
    def receive_messages(self):  
        while True:  
            try:  
                result = self.ws.recv()  
                message = decrypt_data(result)
                self.handle_message(message)  
            except WebSocketConnectionClosedException:  
                print("WebSocket connection closed unexpectedly, attempting to reconnect...")  
                self.reconnect()
                continue  
            except Exception as e:  
                print(f"Error receiving message: {e}")  
                self.reconnect()
                continue  
  
    def reconnect(self):
        """重新连接并重启消息接收"""
        self.close()  # 关闭当前连接
        self.connect()  # 尝试重新连接
        print("Reconnected to WebSocket server.")
  
    def handle_message(self, message):  
        with self._lock:  
            request_id = message.get("request_id")  
            if request_id and request_id in self.pending_requests:  
                # 存储响应到字典中  
                self.response_data[request_id] = message  
                self.pending_requests[request_id].put(message)
            else:
                print(message.get("message")  )
                pass

    def _send_request(self, message):  
        if self.ws is None:  
            print("_send_request failed, websocket not connected")
            return
        
        if not self.ws.connected:  
            print('_send_request failed, websocket not connected')
            return
        
        try:  
            id = -1  
            with self._lock:    
                id = self.request_id    
                self.request_id += 1    
                message["request_id"] = id    
                self.pending_requests[id] = queue.Queue()    
                self.ws.send_binary(encrypt_data(message))  
            
            # 等待响应（这里会阻塞）  
            response = self.pending_requests[id].get(timeout=5)  
            with self._lock:    
                del self.pending_requests[id]  
            
        except Exception as e:  
            print(f"Error in _send_request: {e}")  
            response = None  # 在发生异常时设置response为None  
        
        return response
  
    def API_FindPath(self, start_x, start_y, start_z, end_x, end_y, end_z):  
        request = {  
            "cmd": "FindPath",  
            "start_x": start_x,  
            "start_y": start_y,  
            "start_z": start_z,  
            "end_x": end_x,  
            "end_y": end_y,  
            "end_z": end_z  
        }  
        response_data = self._send_request(request)  
        if response_data is not None:  
            return response_data.get("path", [])  
        else:  
            # 处理 response_data 为 None 的情况  
            return []  # 或者抛出异常，或者返回特定的错误信息
    
    def API_FindRandomPoint(self, cx, cy, cz, min_radius , max_radius):  
        request = {  
            "cmd": "FindRandomPoint",  
            "cx": cx,  
            "cy": cy,  
            "cz": cz,  
            "min_radius": min_radius,
            "max_radius": max_radius
        }  
        response_data = self._send_request(request)  

        if response_data is not None:  
            return response_data.get("point", [])  
        else:  
            # 处理 response_data 为 None 的情况  
            return []  # 或者抛出异常，或者返回特定的错误信息
  
    def start_listening(self):  
        threading.Thread(target=self.receive_messages, daemon=True).start()  

    def start_heartbeat(self):
        threading.Thread(target=self.send_heartbeat, daemon=True).start()
  
    def send_heartbeat(self):
        heartbeat_fail_count = 0
        while True:
            time.sleep(self.heartbeat_interval)
            try:
                self.ws.ping()
                # self._send_request({"cmd": "ping"})
                # print("Heartbeat ping sent.")
                heartbeat_fail_count = 0  # 重置失败计数
            except Exception as e:
                heartbeat_fail_count += 1
                print(f"Failed to send heartbeat (attempt {heartbeat_fail_count}): {e}")
                if heartbeat_fail_count >= self.heartbeat_fail_threshold:
                    print("Heartbeat failed too many times, reconnecting...")
                    self.reconnect()
                    heartbeat_fail_count = 0  # 重置失败计数
  
    def close(self):  
        if self.ws is not None:  
            self.ws.close()  
            print("WebSocket connection closed.")  
  
def start_client(config, client_id):  
    client = MapClient(config)  
    sum = 0  
    while True:  
        response = client.API_FindPath(-466342.53125, 571593.75, -50016.609375, -349447.75, 555150.0625, -45248.66796875)  
        print(f"Client {client_id}: response length {len(response)}, sum: {sum}")  
        sum += 1  
        time.sleep(0.5)  
    client.close()

  
# # 主程序  
# config = {"url": "ws://202.189.9.39:8989/?key=52ad42ad-eee5-47cf-8e40-75391ad9df071"}  
# # config = {"url": "ws://localhost:8989"}  
# clients = []  
# num_clients = 300

# client = MapClient(config)

# while True:
#     response = client.API_FindPath(-466342.53125, 571593.75, -50016.609375, -349447.75, 555150.0625, -45248.66796875)  
#     time.sleep(0.5)
  
# # 创建并启动客户端线程  
# for i in range(num_clients):  
#     client_thread = threading.Thread(target=start_client, args=(config, i))  
#     client_thread.start()  
#     clients.append(client_thread)  
  
# # 等待所有客户端线程完成（在这个例子中，它们实际上不会自行完成，因为内部循环是无限的）  
# # 如果你想要在某个条件满足后停止它们，你需要实现一个停止机制  
# for client in clients:  
#     client.join()  # 这实际上不会返回，因为线程不会自行结束